/**
 * @file FunByMatrixDot.h
 * @author 世慷 (shikang999@126.com)
 * @brief 主要用于矩阵乘法计算系数
 * @version 0.1
 * @date 2025-10-01
 * 
 * @copyright Copyright (c) 2025
 * 
 */

#pragma once
#include "BaseFun.h"
#include "NetLayer.h"
#include "NodeVarForArray.h"
#include "NodeVarForMatrix.h"
#include "NodeVarForVector.h"

namespace Net
{

class FunByMatrixDot : public BaseFun
{
public:
    FunByMatrixDot() : BaseFun(FunType_E::W_dot_x_E, "W * x")
    {

    }

    ~FunByMatrixDot()
    {

    }
    
    void fun(std::shared_ptr<NetLayer> &layer) final
    {
        runFun(vW, layer->vecInputLayer[0], layer->outLayer);
    }

    void diff(std::shared_ptr<NetLayer> &layer) final
    {
        runDiff(vW, layer->vecInputLayer[0], layer->outLayer);
    }

private:
    void runFun(std::shared_ptr<BaseNodeVar> &vW, std::shared_ptr<BaseNodeVar> &x, std::shared_ptr<BaseNodeVar> &y) const
    {
        //这里不检查输入合理性,因为在运行之前需要单独检查
        //1.
        if (vW->getType() == NodeVarType_E::Array_E)
        {
            auto sW = static_cast<NodeVarForArray*>(vW.get());
            auto sx = static_cast<NodeVarForArray*>(x.get());
            auto sy = static_cast<NodeVarForArray*>(y.get());
            for (std::size_t i = sW->vecData.size(); i;)
            {
                --i;
                runFun(sW->vecData[i], sx->vecData[i], sy->vecData[i]);
            }

            return;
        }

        //2.
        {
            auto sW   = static_cast<NodeVarForMatrix*>(vW.get());
            auto sx   = static_cast<NodeVarForVector*>(x.get());
            auto sy   = static_cast<NodeVarForVector*>(y.get());
            sy->value = sW->value * sx->value;
        }
    }

    void runDiff(std::shared_ptr<BaseNodeVar> &vW, std::shared_ptr<BaseNodeVar> &x, std::shared_ptr<BaseNodeVar> &y) const
    {
        //这里不检查输入合理性,因为在运行之前需要单独检查
        //1.
        if (vW->getType() == NodeVarType_E::Array_E)
        {
            auto sW = static_cast<NodeVarForArray*>(vW.get());
            auto sx = static_cast<NodeVarForArray*>(x.get());
            auto sy = static_cast<NodeVarForArray*>(y.get());
            for (std::size_t i = sW->vecData.size(); i;)
            {
                --i;
                runFun(sW->vecData[i], sx->vecData[i], sy->vecData[i]);
            }

            return;
        }

        //2.
        {
            auto sW           = static_cast<NodeVarForMatrix*>(vW.get());
            auto sx           = static_cast<NodeVarForVector*>(x.get());
            auto sy           = static_cast<NodeVarForVector*>(y.get());
            Matrix_t &srcW    = sW->value;
            Matrix_t &diffW   = sW->diffValue;
            Vector_t &srcX    = sx->value;
            Vector_t &diffX   = sx->diffValue;
            auto diffY        = sy->diffValue.data();
            Eigen::Index rows = srcW.rows();
            Eigen::Index cols = srcW.cols();
            for (Eigen::Index i = 0; i != rows; ++i, ++diffY)
            {
                auto ax = srcX.data();
                auto dx = diffX.data();
                for (Eigen::Index j = 0; j != rows; ++j, ++ax, ++dx)
                {
                    //y[i] += srcW[i,j] * srcX[j];
                    diffW(i,j) += (*diffY) * (*ax);
                    (*dx)      += (*diffY) * srcW(i,j);
                }
            }
        }
    }

public:
    std::shared_ptr<BaseNodeVar> vW;//待求解变量, 矩阵运算 vW * x = y
};
}